package llm

import (
	"encoding/json"
	"errors"
	"fmt"
	"refusea-go/utils"
	"time"

	clt "github.com/alibabacloud-go/bailian-20230601/client"
	openapi "github.com/alibabacloud-go/darabonba-openapi/v2/client"
)

/* 阿里百炼 */

var ali_ak = ""
var ali_sk = ""

var bailian_agent_key = ""

const qwen_max_app_id = ""
const qwen_plus_app_id = ""

var bailian_endpoint = "bailian.cn-beijing.aliyuncs.com"

const bailian_url = "https://bailian.aliyuncs.com/v2/app/completions"

var bailian_config = &openapi.Config{
	AccessKeyId:     &ali_ak,
	AccessKeySecret: &ali_sk,
	Endpoint:        &bailian_endpoint,
}

// application id for model
var app_ids = map[string]string{
	MODEL_QWEN_MAX:  qwen_max_app_id,
	MODEL_QWEN_PLUS: qwen_plus_app_id,
}

/* 通义千问, name = [qwen-max, qwen-plus] */
type Qwen struct {
	Name string
	// expire is unix timestamp in sceond
	*Token
}

/**
 * how to access bailian api
 *
 * 1. create ram user, note: the main account may not access the api, i failed.
 * 2. grant right [AliyunSFMFullAccess] to the user
 * 3. add this user in bailian console
 * 4. grant the biz space which your app in it to the user
 *
 **/
func (q *Qwen) RefreshToken() error {

	if q.expire > time.Now().Unix() {
		return nil
	}

	client, err := clt.NewClient(bailian_config)
	if err != nil {
		return err
	}

	request := clt.CreateTokenRequest{AgentKey: &bailian_agent_key}
	result, err := client.CreateToken(&request)
	if err != nil {
		return err
	}

	body := result.Body

	if *body.Success {
		q.token = *body.Data.Token
		q.expire = *body.Data.ExpiredTime - 5

		return nil
	}

	return errors.New(*body.Message)
}
func (q *Qwen) CallApi(system_prompt, user_prompt string) (string, error) {

	rsp, err := client.R().
		SetHeader("Authorization", fmt.Sprintf("Bearer %s", q.token)).
		SetHeader("Content-Type", "application/json; charset=utf-8").
		SetBody(generateQwenReq(system_prompt, user_prompt, app_ids[q.Name])).
		Post(bailian_url)

	if err != nil {
		return "", err
	}

	var qwenRsp QwenRsp
	err = json.Unmarshal(rsp.Body(), &qwenRsp)

	if err != nil {
		return "", err
	}

	if !qwenRsp.Success {
		return "", errors.New(qwenRsp.Message)
	}

	return qwenRsp.Data.Text, nil
}

func generateQwenReq(system_prompt, user_prompt string, app_id string) map[string]interface{} {
	return map[string]interface{}{
		"RequestId": utils.Next(time.Now().UnixMilli()),
		"AppId":     app_id,
		"Messages": []map[string]string{
			{
				"Role":    "system",
				"Content": system_prompt,
			}, {
				"Role":    "user",
				"Content": user_prompt,
			},
		},
	}
}

type QwenChoiceMessage struct {
	Role    string `json:"Role"`
	Content string `json:"Content"`
}

type QwenChoice struct {
	FinishReason string            `json:"FinishReason"`
	Message      QwenChoiceMessage `json:"Message"`
}

type QwenUsage struct {
	InputTokens  int `json:"InputTokens"`
	OutputTokens int `json:"OutputTokens"`
}

type QwenData struct {
	ResponseId string       `json:"ResponseId"`
	Text       string       `json:"Text"`
	Choices    []QwenChoice `json:"Choices"`
	Usage      []QwenUsage  `json:"Usage"`
}

type QwenRsp struct {
	Success   bool     `json:"Success"`
	Code      string   `json:"Code"`
	Message   string   `json:"Message"`
	RequestId string   `json:"RequestId"`
	Data      QwenData `json:"Data"`
}
